# 要添加一个新单元，输入 '# %%'
# 要添加一个新的标记单元，输入 '# %% [markdown]'
# 我发现写到最后，还是学弟写的那种样子。好麻烦啊，感觉应该一天就能写完，但自己在硬拖。


# %%
import torch
import copy
from torch import nn, optim
from dnn_model import Net
from torch.utils import data 


# 训练相关的参数本来应该是服务器发送过来的，包括模型、参数啥的。训练数据应该是本地的。
# 这里为了方便，只模拟联邦学习的效果和逻辑，不模拟联邦学习的真是过程。
# 1. 数据由服务器分成独立同分布，或者非独立同分布。然后分配给各个客户端。而非客户端在本地读取。
# 2. 训练的模型、参数在客户端配置。
# 3. 客户端完成训练过程、服务端完成聚合过程。客户端和服务端只发送和传输模型和优化器的权重参数。完成聚合和分配的通信逻辑。

class FLClient:
    # 接收参数，完成训练。
    def __init__(self,data):
        self.data = data
        self.model = Net()
        self.loss_func = nn.CrossEntropyLoss()
        self.optimizer = optim.Adam(params=self.model.parameters(),lr=0.001)
        self.epochs =1
        self.batch_size=256
        self.loss_prox = True
    

    def load_model(self,model_state_dict,optimizer_state_dict):
        self.model.load_state_dict(model_state_dict)
        self.optimizer.load_state_dict(optimizer_state_dict)
    

    # 用来测试训练的结果返回准确率
    def test(self):
        self.model.eval()
        x,y = self.data[:,0:470],self.data[:,470]
        # 将数据转换成tensor
        x_test_tensor = torch.FloatTensor(x)
        y_test_tensor = torch.LongTensor(y)
        test_ds = data.TensorDataset(x_test_tensor,y_test_tensor)
        test_dl = data.DataLoader(test_ds,batch_size=self.batch_size,shuffle=True)

        correct = 0
        total =0
        
        for features,labels in test_dl:
            outputs = self.model(features)
            # print(outputs.shape)
            _,predicted = torch.max(outputs,1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

        accuracy = 100*correct/total
        return accuracy


    # 用来训练当前的本地模型
    def train(self):
        self.model.train()
        x,y = self.data[:,0:470],self.data[:,470]
           # 将数据转换成tensor
        x_train_tensor = torch.FloatTensor(x)
        y_train_tensor = torch.LongTensor(y)
        train_ds = data.TensorDataset(x_train_tensor,y_train_tensor)
        train_dl = data.DataLoader(train_ds,batch_size=self.batch_size,shuffle=True)

        for epoch in range(self.epochs):
            for xb,yb in train_dl:
                self.optimizer.zero_grad()
                pred = self.model(xb)
                loss = self.loss_func(pred,yb.long())/len(train_dl)

                loss.backward()
                self.optimizer.step()
                self.optimizer.zero_grad()

        # print(self.optimizer.state_dict())

    # 返回当前的模型
    def get_model_state_dict(self):
        return self.model.state_dict()

    # 返回当前的优化器
    def get_optim_state_dict(self):
        return self.optimizer.state_dict()
 